# -*- coding: utf-8 -*-
"""
项目名称：cotton_test
文件名称：train.py
创建时间：2025-06-08 20:47:50

系统用户：wyl
作　　者：無以菱
联系邮箱：huangjing2001.guet@qq.com
功能描述：棉花叶片疾病分类模型训练脚本
- 本脚本用于训练两种深度学习模型（BP神经网络和卷积神经网络）进行棉花叶片的健康/患病分类。
- 脚本会自动加载数据、训练模型、评估性能并保存训练好的模型和结果图表。
"""


import os
import sys
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, random_split
import matplotlib.pyplot as plt
import numpy as np
from tqdm import tqdm  # 进度条库

# 导入自定义模块
from utils.transform import transform  # 图像预处理转换
from utils.dataset import CottonLeafDataset  # 数据集类
from utils.accuracy import accuracy  # 准确率计算函数
from model import SimpleMLP, CustomCNN  # 模型定义

# 设置随机种子，确保结果可复现（相同的初始条件下得到相同的结果）
torch.manual_seed(42)  # PyTorch随机种子
np.random.seed(42)     # NumPy随机种子

# 检查是否有可用的GPU，如果有则使用GPU，否则使用CPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(f"使用设备: {device}")  # 输出当前使用的计算设备

# 数据集路径
data_dir = "data"

# 检查数据目录是否存在
if not os.path.exists(data_dir):
    print(f"错误: 数据目录 '{data_dir}' 不存在!")
    sys.exit(1)

# 列出数据目录中的内容
print(f"数据目录 '{data_dir}' 中的内容:")
for item in os.listdir(data_dir):
    item_path = os.path.join(data_dir, item)
    if os.path.isdir(item_path):
        print(f"  - 文件夹: {item} (包含 {len(os.listdir(item_path))} 个文件)")
    else:
        print(f"  - 文件: {item}")

# 创建数据集
try:
    dataset = CottonLeafDataset(root_dir=data_dir, transform=transform)
except Exception as e:
    print(f"创建数据集时出错: {str(e)}")
    sys.exit(1)

# 检查数据集是否为空
if len(dataset) == 0:
    print("错误: 数据集为空! 请检查数据目录结构和文件格式。")
    sys.exit(1)

# 获取类别数量
num_classes = len(dataset.label_map)
print(f"数据集包含 {num_classes} 个类别: {dataset.label_map}")
print(f"总样本数: {len(dataset)}")

# 划分训练集和测试集 (80% 训练, 20% 测试)
train_size = int(0.8 * len(dataset))
test_size = len(dataset) - train_size
train_dataset, test_dataset = random_split(dataset, [train_size, test_size])

print(f"训练集样本数: {len(train_dataset)}")
print(f"测试集样本数: {len(test_dataset)}")

# 创建数据加载器
batch_size = 32
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

# 创建模型实例
def train_model(model_type="cnn", num_epochs=10):
    """
    训练模型
    
    参数:
        model_type (str): 模型类型，"mlp" 或 "cnn"
        num_epochs (int): 训练轮数
        
    返回:
        model: 训练好的模型
        train_acc_history: 训练准确率历史
        test_acc_history: 测试准确率历史
    """
    if model_type == "mlp":
        model = SimpleMLP(num_classes=num_classes)
        model_name = "BP神经网络"
        save_path = "model/mlp_model.pth"
    else:
        model = CustomCNN(num_classes=num_classes)
        model_name = "卷积神经网络"
        save_path = "model/cnn_model.pth"
    
    model = model.to(device)
    print(f"使用 {model_name} 模型")
    
    # 定义损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001, weight_decay=5e-4)
    
    # 学习率调度器
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, mode='min', factor=0.5, patience=5, min_lr=0.000001
    )
    
    # 记录训练过程
    train_acc_history = []
    test_acc_history = []
    best_test_acc = 0.0
    best_model_state = None
    
    # 训练循环
    for epoch in range(num_epochs):
        # 训练阶段
        model.train()
        train_loss = 0.0
        train_correct = 0
        train_total = 0
        
        # 使用tqdm显示进度条
        train_pbar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{num_epochs} [Train]")
        for data, target in train_pbar:
            data, target = data.to(device), target.to(device)
            
            # 前向传播
            optimizer.zero_grad()
            outputs = model(data)
            loss = criterion(outputs, target)
            
            # 反向传播和优化
            loss.backward()
            optimizer.step()
            
            # 统计训练指标
            train_loss += loss.item()
            right, total = accuracy(outputs, target)
            train_correct += right.item()
            train_total += total
            
            # 更新进度条
            train_pbar.set_postfix({'loss': loss.item(), 'acc': 100.0 * right.item() / total})
        
        # 计算平均训练损失和准确率
        train_loss = train_loss / len(train_loader)
        train_acc = 100.0 * train_correct / train_total
        train_acc_history.append(train_acc)
        
        # 测试阶段
        model.eval()
        test_loss = 0.0
        test_correct = 0
        test_total = 0
        
        with torch.no_grad():
            test_pbar = tqdm(test_loader, desc=f"Epoch {epoch+1}/{num_epochs} [Test]")
            for data, target in test_pbar:
                data, target = data.to(device), target.to(device)
                
                # 前向传播
                outputs = model(data)
                loss = criterion(outputs, target)
                
                # 统计测试指标
                test_loss += loss.item()
                right, total = accuracy(outputs, target)
                test_correct += right.item()
                test_total += total
                
                # 更新进度条
                test_pbar.set_postfix({'loss': loss.item(), 'acc': 100.0 * right.item() / total})
        
        # 计算平均测试损失和准确率
        test_loss = test_loss / len(test_loader)
        test_acc = 100.0 * test_correct / test_total
        test_acc_history.append(test_acc)
        
        # 保存最佳模型
        if test_acc > best_test_acc:
            best_test_acc = test_acc
            best_model_state = model.state_dict()
        
        # 调整学习率
        scheduler.step(test_loss)
        
        # 打印训练信息
        print(f"Epoch {epoch+1}/{num_epochs}: "
              f"Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.2f}%, "
              f"Test Loss: {test_loss:.4f}, Test Acc: {test_acc:.2f}%, "
              f"Best Test Acc: {best_test_acc:.2f}%")
    
    # 加载最佳模型
    if best_model_state:
        model.load_state_dict(best_model_state)
        # 保存模型
        torch.save(model.state_dict(), save_path)
        print(f"模型已保存到 {save_path}")
    
    return model, train_acc_history, test_acc_history

def main():
    """
    主函数：训练两种模型并可视化结果
    
    流程:
    1. 训练MLP模型（多层感知机）
    2. 训练CNN模型（卷积神经网络）
    3. 绘制并保存训练曲线图表
    4. 保存训练好的模型
    """
    try:
        # 训练MLP模型（多层感知机/BP神经网络）
        mlp_model, mlp_train_acc, mlp_test_acc = train_model(model_type="mlp", num_epochs=10)
        print("MLP模型训练完成!")

        # 训练CNN模型（卷积神经网络）
        cnn_model, cnn_train_acc, cnn_test_acc = train_model(model_type="cnn", num_epochs=10)
        print("CNN模型训练完成!")

        # 绘制训练曲线，比较两种模型的性能
        plt.figure(figsize=(12, 5))  # 创建图表，设置大小

        # 绘制训练和测试准确率曲线
        plt.subplot(1, 2, 1)  # 创建子图
        plt.plot(mlp_train_acc, 'b-', label='MLP Train Acc')  # MLP训练准确率（蓝色实线）
        plt.plot(mlp_test_acc, 'b--', label='MLP Test Acc')   # MLP测试准确率（蓝色虚线）
        plt.plot(cnn_train_acc, 'r-', label='CNN Train Acc')  # CNN训练准确率（红色实线）
        plt.plot(cnn_test_acc, 'r--', label='CNN Test Acc')   # CNN测试准确率（红色虚线）
        plt.title('Training and Testing Accuracy')  # 图表标题
        plt.xlabel('Epoch')  # x轴标签
        plt.ylabel('Accuracy (%)')  # y轴标签
        plt.legend()  # 显示图例
        plt.grid(True)  # 显示网格

        # 保存图表到结果目录
        plt.tight_layout()  # 调整布局
        plt.savefig('results/training_curves.png')  # 保存图表
        plt.show()  # 显示图表（如果在支持图形界面的环境中）

        print("训练完成!")

    except Exception as e:
        # 捕获并打印训练过程中的错误
        print(f"训练过程中出错: {str(e)}")
        sys.exit(1)  # 异常退出

if __name__ == "__main__":
    # 创建结果目录
    os.makedirs("results", exist_ok=True)
    main() 